"""
Example calculations of photoionization rate. 
"""

import sys
import numpy as np
import matplotlib.pyplot as plt

from scipy.constants import codata


# set up constants
#-------------------------------------------------------------------------

const = codata.physical_constants

Rbohr = const['Bohr radius'][0] * 100        # Bohr radius [cm]
alpha = const['fine-structure constant'][0]  # fine structure constant

# http://adsabs.harvard.edu/abs/2009RvMP...81.1405M (Eq. 47)
# H photoionization cross-section at nu_th 

A0 = (2**9 * np.pi) / (3 * np.exp(1)**4) * alpha * np.pi * Rbohr**2



# define functions
#-------------------------------------------------------------------------

def regularize_input( ryd_in ):
    """ add epsilon to input ryd values equal to 1 """

    if isinstance(ryd_in,float):
        ryd = ryd_in
        if ryd_in == 1:
            ryd = 1 + sys.float_info.epsilon
    else:
        ryd = ryd_in.copy()
        indx = np.where( ryd_in == 1 )
        ryd[indx] = 1 + sys.float_info.epsilon

    return ryd


def regularize_output( ryd_in, sigma ):
    """ set 0 cross-section for input ryd values below 1 """

    if isinstance(ryd_in,float):
        if ryd_in < 1:
            sigma = 0
    else:
        indx = np.where( ryd_in < 1 )
        sigma[indx] = 0

    return sigma


def h1_photo_sigma_analytic( ryd_in ):
    """ neutral hydrogen photoionization cross-section - analytic 
    http://adsabs.harvard.edu/abs/2009RvMP...81.1405M (Eq. 47) """

    # regularize input
    #--------------------------------------------
    ryd = regularize_input(ryd_in)
    
    # calculate cross-section
    #--------------------------------------------
    eps = np.sqrt(ryd - 1)
    ieps = 1/eps
    efac = 4 - ( 4 * np.arctan(eps) ) * ieps
    num = np.exp(efac)
    den = 1 - np.exp(-2 * np.pi * ieps)        
    sigma = A0 * (ryd)**(-4) * num / den

    # regularize output
    #--------------------------------------------
    sigma = regularize_output(ryd_in, sigma)
    return sigma


def h1_photo_sigma_verner( ryd_in ):
    """ neutral hydrogen photoionization cross-section - verner 
    http://adsabs.harvard.edu/abs/1996ApJ...465..487V (Table 1) """

    # regularize input
    #--------------------------------------------
    ryd = regularize_input(ryd_in)

    # calculate cross-section
    #--------------------------------------------
    Eth = 13.6
    Emax = 5.0e4
    E0 = 4.298e-1
    sig0 = 5.475e4
    ya = 3.288e1
    P = 2.963
    yw = 0.0
    y0 = 0.0
    y1 = 0.0
        
    eV = ryd * Eth
    x = eV / E0 - y0
    y = x
        
    fac = ( 1 + np.sqrt(y/ya) )**(-P)
    sigma = sig0 * (x-1)**2 * y**(0.5 * P - 5.5) * fac 
    sigma = sigma * 1.0e-18

    # regularize output
    #--------------------------------------------
    sigma = regularize_output(ryd_in, sigma)
    return sigma



def h1_photo_sigma_osterbrock( ryd_in ):
    """ neutral hydrogen photoionization cross-section - osterbrock 
    http://adsabs.harvard.edu/abs/2009RvMP...81.1405M (Eq. 48) """

    # regularize input
    #--------------------------------------------
    ryd = regularize_input(ryd_in)
    
    # calculate cross-section
    #--------------------------------------------
    aT = 6.30e-18
    beta = 1.34
    s = 2.99
    sigma = aT * ( beta * ryd**(-s) + (1-beta) * ryd**(-s-1) )
    
    # regularize output
    #--------------------------------------------
    sigma = regularize_output(ryd_in, sigma)
    return sigma


def h1_photo_sigma_powerlaw( ryd_in ):
    """ neutral hydrogen photoionization cross-section - single powerlaw """

    # regularize input
    #--------------------------------------------
    ryd = regularize_input(ryd_in)
    
    # calculate cross-section
    #--------------------------------------------
    sigma = A0 * ryd**(-3)
    
    # regularize output
    #--------------------------------------------
    sigma = regularize_output(ryd_in, sigma)
    return sigma



def h1_photo_sigma(ryd, eval="verner"):
    """ Hydrogen photo-ionization cross section. """ 

    # choose evaluation method and return
    #--------------------------------------------

    if eval == "verner":

        sigma = h1_photo_sigma_verner( ryd )
        return sigma

    elif eval == "osterbrock":

        sigma = h1_photo_sigma_osterbrock( ryd )
        return sigma

    elif eval == "analytic":

        sigma = h1_photo_sigma_analytic( ryd )
        return sigma

    elif eval == "powerlaw":
        sigma = h1_photo_sigma_powerlaw( ryd )
        return sigma

    else:
        print "eval method, " + eval + ", not recognized"
        sys.exit(1)






def plot_sigmas():


    ryd = np.linspace( 0.1, 30.0, 50000 )

    sigma_analytic = h1_photo_sigma_analytic( ryd )
    sigma_verner = h1_photo_sigma_verner( ryd )
    sigma_osterbrock = h1_photo_sigma_osterbrock( ryd )
    sigma_powerlaw = h1_photo_sigma_powerlaw( ryd )

    #plt.plot( ryd, np.log10(sigma_analytic) )
    #plt.plot( ryd, np.log10(sigma_verner) )
    #plt.plot( ryd, np.log10(sigma_osterbrock) )


    plt.plot( ryd, sigma_verner/sigma_analytic,
              color='blue')

    plt.plot( ryd, sigma_osterbrock/sigma_analytic, 
              color='green')

    plt.plot( ryd, sigma_powerlaw/sigma_analytic, 
              color='red')
